// Copyright (C) 2009 Joshua Emele
// and all other respective authors
//
// jmem is free software; you can redistribute it and/or
// modify it under the terms of the GNU General Public
// License as published by the Free Software Foundation; either
// version 3 of the License, or (at your option) any later version.
//
// jmem is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  
//
// You should have received a copy of the GNU General Public
// License along with jmem; if not, write to the Free
// Software Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA
// 02111-1307 USA.  
// 
// Parts derived from binutils/addr2line.c, which in turn was 
// derived from objdump.c and nm.c by Ulrich.Lauther@mchp.siemens.de
//
// If you're using this, you probably want to run with
// MALLOC_CHECK_=0 or MALLOC_CHECK_=1
//
#ifndef _jmem_h
#define _jmem_h
#define JMEM_LOGV(format, ...) \
  fprintf(stderr, "%s:%d: %s: " format "\n", basename(__FILE__), \
      __LINE__, __PRETTY_FUNCTION__,  __VA_ARGS__)
#define JMEM_LOG(message, ...) JMEM_LOGV(message, 0)

#include <bfd.h>
#include <dirent.h> 
#include <errno.h>
#include <execinfo.h>
#include <fcntl.h>
#include <libiberty.h>
#include <linux/types.h>
#include <malloc.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/stat.h> 
#include <sys/types.h>
#include <unistd.h>

// taken from demangle.h
#define DMGL_NO_OPTS	         0	// For readability... 
#define DMGL_PARAMS      (1 << 0)	// Include function args 
#define DMGL_ANSI        (1 << 1)	// Include const, volatile, etc
#define DMGL_JAVA	       (1 << 2)	// Demangle as Java rather than C++
#define DMGL_VERBOSE	   (1 << 3)	// Include implementation details
#define DMGL_TYPES	     (1 << 4)	// Also try to demangle type encodings
#define DMGL_RET_POSTFIX (1 << 5) // Print function return types (when
                                  // present) after function signature 
#define DMGL_AUTO	 (1 << 8)
#define DMGL_GNU	 (1 << 9)
#define DMGL_LUCID (1 << 10)
#define DMGL_ARM	 (1 << 11)
#define DMGL_HP 	 (1 << 12)      // For the HP aCC compiler;
                                  // same as ARM except for
                                  // template arguments, etc
#define DMGL_EDG	  (1 << 13)
#define DMGL_GNU_V3	(1 << 14)
#define DMGL_GNAT	  (1 << 15)

struct jmem_header_t {
  ssize_t count;
  size_t size;
};

struct jmem_find_addr_context_t {
  bfd_vma pc;
  asymbol** syms;
  const char* filename;
  const char* functionname;
  unsigned int line;
  bfd_boolean found;
};

static void jmem_atexit(void);
static void jmem_init(void);
static void jmem_find_address_in_section(bfd*, asection*, void*);
static void jmem_free(void*, const void*);
static void* jmem_malloc(size_t, const void*);
static void jmem_record_malloc(size_t, void*);
static int jmem_record_free(struct jmem_header_t*, void*);
static asymbol** jmem_slurp_symtab(bfd*);
static int jmem_temporary_fd();
static void jmem_scan_trace_line(char* line, char** trace_file, 
  char** trace_func, char** trace_addr);
static int jmem_trace_line(struct jmem_find_addr_context_t*, const char*, 
  const char*);
static int jmem_trace_record(FILE*, void**, uint32_t);

void (*__malloc_initialize_hook)(void) = jmem_init;
void *(*jmem_old_malloc_hook)(size_t, const void *) = 0;
void (*jmem_old_free_hook)(void*, const void *) = 0;

//*****************************************************************************
// Read in the symbol table.
static asymbol** jmem_slurp_symtab(bfd* abfd) 
{
  // check for symbols
  if (!(bfd_get_file_flags(abfd) & HAS_SYMS)) {
    JMEM_LOGV("No symbols found for %s", bfd_get_filename(abfd));
    return 0;
  }

  // read symbols and validate
  asymbol** syms = 0;
  unsigned int size = 0;
  long symcount = bfd_read_minisymbols(abfd, FALSE, (void**)&syms, &size);
  if (!symcount) {
    symcount = bfd_read_minisymbols(abfd, TRUE/*dynamic*/, (void**)&syms, 
      &size);
  }
  if (symcount < 0) {
    bfd_error_type error = bfd_get_error();
    JMEM_LOGV("Error finding symbols for %s: %s",
      bfd_get_filename(abfd), bfd_errmsg(error));
  }

  return syms;
}

//*****************************************************************************
// Look for an address in a section.  This is called via bfd_map_over_sections. 
static void jmem_find_address_in_section(bfd* abfd, asection* section, void* data)
{
  if (!data) {
    return;
  }
  struct jmem_find_addr_context_t* context = 
    (struct jmem_find_addr_context_t*)data;
  if (context->found) {
    return;
  }
   
  if (0 == (bfd_get_section_flags(abfd, section) & SEC_ALLOC)) {
    return;
  }

  bfd_vma vma = bfd_get_section_vma(abfd, section);
  if (context->pc < vma) {
    return;
  }

  bfd_size_type size = bfd_get_section_size(section);
  if (context->pc >= vma + size) {
    return;
  }

  context->found = bfd_find_nearest_line(abfd, section, context->syms, 
    context->pc - vma, &context->filename, &context->functionname, 
    &context->line);
// XXX extra goodies
// alloc = bfd_demangle(abfd, name, DMGL_ANSI | DMGL_PARAMS);
// found = bfd_find_inliner_info(abfd, &filename, &functionname, &line);
}

//*****************************************************************************
static int jmem_trace_line(struct jmem_find_addr_context_t* context, 
  const char* file, const char* address)
{
  // verify the file exists and has size
  struct stat st = {0};
  if (-1 == stat(file, &st) || st.st_size < 1)  {
    int error = errno;
    JMEM_LOGV("state failed for %s with error: %s (size %d)",
      file, strerror(error), st.st_size);
    return error;
  }

  // open the bfd and validate
  bfd* abfd = bfd_openr(file, 0);
  if (!abfd) {
    bfd_error_type error = bfd_get_error();
    JMEM_LOGV("bfd_oper failed for file %s and with error: %s",
      file, bfd_errmsg(error));
    return error;
  }
  
  // verify the format checks out ...
  if (bfd_check_format(abfd, bfd_archive)) {
    bfd_error_type error = bfd_get_error();
    JMEM_LOGV("bfd_check_format failed with error: %s", 
      bfd_errmsg(error));
    bfd_close(abfd);
    return error;
  }
  
  // XXX necessary?
  char** matching = 0;
  if (!bfd_check_format_matches(abfd, bfd_object, &matching)) {
    bfd_error_type error = bfd_get_error();
    JMEM_LOGV("bfd_check_format_matches failed with error: %s",
      bfd_errmsg(error));
    bfd_close(abfd);
    return error;
  }
      
  // XXX caching/performance?
  // this collapses as much as we can ... we want this to be a one
  // size fits all function ...
  context->syms = jmem_slurp_symtab(abfd);
  if (!context->syms) {
    JMEM_LOGV("jmem_slurp_symtab failed for %s", file);
    return EINVAL;
  }

  // release the hounds ...
  context->pc = bfd_scan_vma(address, 0, 16);
  bfd_map_over_sections(abfd, jmem_find_address_in_section, context);
  free(context->syms);
  context->syms = 0;

  // close and all is well
  bfd_close(abfd);
  return 0;
}

//*****************************************************************************
static void jmem_init(void)
{
  // initialize bfd
  bfd_init();
  char* target = getenv("JMEM_BFD_TARGET");
  if (!target || !strlen(target)) {
    target = (char*)"i386-redhat-linux-gnu";
  }
  JMEM_LOGV("Assuming target: %s", target);
  if (!bfd_set_default_target(target)) {
    bfd_error_type error = bfd_get_error();
    JMEM_LOGV("Failed to set BFD default target to %s with error: %s",
      target, bfd_errmsg(error));
    exit(error);
  }

  // allocate record directory
  char path[32] = {0};
  snprintf(path, sizeof(path), "/tmp/%d", getpid());
  if (-1 == mkdir(path, 0700)) {
    int error = errno;
    JMEM_LOGV("failed to create %s with error %s",
        path, strerror(error));
    exit(error);
  }
  JMEM_LOGV("Writing alloc records to %s", path);
  
  // register an atexit handler
  if (atexit(jmem_atexit)) {
    JMEM_LOG("Failed to register atexit handler");
    exit(EINVAL);
  }
  
  // save old hooks
  jmem_old_malloc_hook = __malloc_hook;
  jmem_old_free_hook = __free_hook;

  // and set our hooks
  __malloc_hook = jmem_malloc;
  __free_hook = jmem_free;
}

//*****************************************************************************
static void jmem_atexit(void)
{
  // restore old hooks
  __malloc_hook = jmem_old_malloc_hook;
  __free_hook = jmem_old_free_hook;

  // open the record directory and validate
  char path[32] = {0};
  snprintf(path, sizeof(path), "/tmp/%d", getpid());
  DIR* records = opendir(path);
  if (!records) {
    int error = errno;
    JMEM_LOGV("opendir failed for %s with error: %s",
      strerror(error));
    return;
  }
   
  /* READDIR(3)
  Since POSIX.1 does not specify the size of the d_name field, and other 
  non-standard fields may precede that field within the dirent structure, 
  portable applications that use readdir_r() should allocate the
  buffer whose address is passed in entry as follows:

     len = offsetof(struct dirent, d_name) +
               pathconf(dirpath, _PC_NAME_MAX) + 1
     entryp = malloc(len);
  */
  union {
    struct dirent entry;
    char padding[offsetof(struct dirent, d_name) + NAME_MAX + 1];
  } buffer = {0};

  // for each entry
  int status = 0;
  struct dirent* entry = 0;
  while (0 == (status = readdir_r(records, &buffer.entry, &entry))) {

    // are we done yet?
    if (!entry) {
      break;
    }

    // skip non-files
    if (DT_REG != entry->d_type) {
      continue;
    }

    // create a path to the record, open, and validate
    snprintf(path, sizeof(path), "/tmp/%d/%s", getpid(), entry->d_name);
    int fd = open(path, O_RDWR);
    if (-1 == fd) {
      int error = errno;
      JMEM_LOG("Failed to open %s with error: %s",
          path, strerror(error));
      unlink(path);
      continue;
    }
    
    // read the header and validate
    struct jmem_header_t header;
    if (sizeof(header) != read(fd, &header, sizeof(header))) {
      int error = errno;
      JMEM_LOGV("Failed to read header for (%s) in %s: %s",
          entry->d_name, path, strerror(error));
      close(fd);
      unlink(path);
      continue;
    }

    // if the resource has been freed, we don't care
    if (header.count <= 0) {
      close(fd);
      unlink(path);
      continue;
    }

    // otherwise, we have a potential leak ...
    JMEM_LOGV("potential leak (%s) size %d count %d",
        entry->d_name, header.size, header.count);
    
    // get the allocation trace 
    char trace[2048] = {0};
    ssize_t bytes = read(fd, trace, sizeof(trace));
    if (bytes < 0) {
      int error = errno;
      JMEM_LOGV("failed to read backtrace for (%s): %s",
          entry->d_name, strerror(error));
    }

    // remove the entry
    close(fd);
    unlink(path);

    // print the backtraces
    JMEM_LOGV("allocation backtrace:\n"
      "%s\ndeallocation backtrace:", trace);
    uint32_t trace_size = 10;
    void *array[trace_size];
    trace_size = backtrace(array, trace_size);
    if (jmem_trace_record(stderr, array, trace_size)) {
      // XXX panic?
      JMEM_LOG("Failed to create a trace record");
    }
    JMEM_LOG("end backtrace\n\n");
  }

  // close the directory and validate
  snprintf(path, sizeof(path), "/tmp/%d", getpid());
  if (-1 == closedir(records)) {
    int error = errno;
    JMEM_LOGV("failed to close record directory %s with error: %s",
      path, strerror(error));
  }

  // remove the directory and validate
  if (-1 == rmdir(path)) {
    int error = errno;
    JMEM_LOGV("failed to remove record directory %s with error: %s",
      path, strerror(error));
  }
}

//*****************************************************************************
static int jmem_temporary_fd() 
{
  char path[32] = {0};
  snprintf(path, sizeof(path), "/tmp/%d/tmp.XXXXXX", getpid());
  int fd = mkstemp(path);
  if (-1 == fd) {
    int error = errno;
    JMEM_LOGV("Failed to open %s with error: %s",
      path, strerror(error));
    return -1;
  }
  unlink(path);
  return fd;
}

//*****************************************************************************
static void jmem_scan_trace_line(char* line, char** trace_file, 
  char** trace_func, char** trace_addr)
{
  // convert to bin/fn/addr
  // the trace will take one the following format forms:
  *trace_file = line;
  *trace_func = 0;
  *trace_addr = 0;

  // i.  file(fn+offset)[addr]
  char* mark = strchr(line, '(');
  if (mark) {
    
    // terminate trace_file
    *mark = 0;

    // locate and terminate trace_func
    *trace_func = ++mark;
    mark = strchr(++mark, ')');
    *mark = 0;

    // locate and terminate trace_addr
    mark = strchr(++mark, '[');
    *trace_addr = ++mark;
    mark = strchr(++mark, ']');
    *mark = 0;

  // ii. file[addr]
  } else {

    // locate and terminate trace_file
    mark = strchr(line, '[');
    *mark = 0;

    // locate and terminate trace addr
    *trace_addr = ++mark;
    mark = strchr(++mark, ']');
    *mark = 0;
  }
}


//*****************************************************************************
static int jmem_trace_record(FILE* record, void** array, uint32_t trace_size)
{
  // open a temporary file to write the backtrace out
  int fd = jmem_temporary_fd();
  if (-1 == fd) {
    JMEM_LOG("Failed to create a temporary file");
    return EBADF;
  }
  backtrace_symbols_fd(array, trace_size, fd);

  // seek and validate
  if (-1 == lseek(fd, 0, SEEK_SET)) {
    int error = errno;
    JMEM_LOGV("Failed to seek on temporary file with error: %s", 
      strerror(error));
    close(fd);
    return error;
  }

  // open a buffered file stream
  FILE* file = fdopen(fd, "r"); 
  if (!file) {
    int error = errno;
    JMEM_LOGV("Failed to open file stream on temporary file with error: %s",
      strerror(error));
    close(fd);
    return error;
  }

  // for each line of backtrace
  int level = 0;
  char line[1024] = {0};
  while (fgets(line, sizeof(line), file)) {

    // ferret out the useful bits
    char* trace_file = 0;
    char* trace_func = 0;
    char* trace_addr = 0;
    jmem_scan_trace_line(line, &trace_file, &trace_func, &trace_addr);

    // XXX write the translated file, line info to the alloc record
    // XXX configurable and default established in jmem_init ... 
    // we don't need to do this everytime
    struct jmem_find_addr_context_t context = {0};
    if (jmem_trace_line(&context, trace_file, trace_addr)) {
      JMEM_LOGV("Failed to trace %s:%s",
        trace_file, trace_addr);
      continue;
    }
    
    // write a record like gdb:
    // (gdb) where
    // #0  jmem_atexit () at ../include/jmem.h:250
    // #1  0x00b9c949 in exit () from /lib/libc.so.6
    // #2  0x00b846ed in __libc_start_main () from /lib/libc.so.6
    // #3  0x08049a01 in _start ()
    fprintf(record, "#%d ", level++);
    if (context.functionname && context.filename) {
      fprintf(record, "%s () at %s:%d\n", 
        context.functionname, 
        context.filename, 
        context.line);
    } else {
      fprintf(record, "%s in %s from %s\n",
        trace_addr, trace_func, trace_file);
    }
  }

  // close and all is well
  fclose(file);
  return 0;
}

//*****************************************************************************
// XXX this is slower than you can imagine ...
// XXX if your code relies on timing, this won't help ....
// XXX consider another way to do this faster
static void jmem_record_malloc(size_t size, void* addr) 
{
  // do not use the heap from here on out...
  // create a record path
  char path[32] = {0};
  snprintf(path, sizeof(path), "/tmp/%d/%p", getpid(), addr);

  // check the record for length
  struct stat st = {0};
  if (-1 != stat(path, &st)) {

    // if the record has length
    if (st.st_size) {

      // open the file and verify count == 0
      struct jmem_header_t header = {0};
      int fd = open(path, O_RDONLY, 0400);
      if (-1 != fd) {
        read(fd, &header, sizeof(header));
        close(fd);
      }
        
      // if the record is valid and count != 0, we're going to truncate
      // which means alloc reused an address without free'ing 
      // or someone free'd more than once ... 
      if (header.size && header.count != 0) {
        JMEM_LOGV("Truncating record for "
            "address %p size %d count %d",
            addr, header.size, header.count);
      }
    }
  }
  
  // open and truncate the record
  int fd = creat(path, 0600);
  if (-1 == fd) {
    int error = errno;
    JMEM_LOGV("Failed to open %s with error %s\n",
        error, strerror(error));
    exit(error);
  }

  // write the header
  struct jmem_header_t header = { 1, size };
  if (sizeof(header) != write(fd, &header, sizeof(header))) {
    int error = errno;
    JMEM_LOGV("Failed to write malloc header for "
        "%p to %s with error: %s",
        addr, path, strerror(error));
    exit(error);
  }

  // create the trace data
  uint32_t trace_size = 10;
  void *array[trace_size];
  trace_size = backtrace(array, trace_size);

  // open a file stream for the trace data
  FILE* record = fdopen(fd, "a");
  if (!record) {
    int error = errno;
    exit(error);
  }

  if (jmem_trace_record(record, array, trace_size)) {
    // XXX panic?
    JMEM_LOG("Failed to create a trace record");
  }
    
  fclose(record);
}

//*****************************************************************************
static void* jmem_malloc(size_t size, const void *caller)
{
  // restore old hooks
  __malloc_hook = jmem_old_malloc_hook;
  __free_hook = jmem_old_free_hook;

  // malloc ...
  void* addr = malloc(size);

  // _try_ not to use the heap from here on out...
  // write out the record header
  // and the backtrace ...
  jmem_record_malloc(size, addr);
 
  // save hooks and restore
  jmem_old_malloc_hook = __malloc_hook;
  jmem_old_free_hook = __free_hook;
  __malloc_hook = jmem_malloc;
  __free_hook = jmem_free;

  // thank you, come again
  return addr;
}

//*****************************************************************************
static int jmem_record_free(struct jmem_header_t* header, void* addr) 
{
  // do not use the heap from here on out...
  // lookup the entry for the address
  char path[32] = {0};
  snprintf(path, sizeof(path), "/tmp/%d/%p", getpid(), addr);

  // verify the record exists and is valid
  struct stat st = {0};
  if (-1 == stat(path, &st)) { 
    int error = errno;
    JMEM_LOGV("Invalid malloc record for %p in %s: %s",
        addr, path, strerror(error));
    return error;
  }
  if (st.st_size < sizeof(struct jmem_header_t)) {
    JMEM_LOGV("Invalid malloc record for %p in %s: header too small (%d)",
        addr, path, st.st_size);
    unlink(path);
    return EINVAL;
  }

  // open the record and validate
  int fd = open(path, O_RDWR);
  if (-1 == fd) {
    int error = errno;
    JMEM_LOGV("Failed to open %s with error: %s",
        path, strerror(error));
    unlink(path);
    return error;
  }
  
  // read the header and validate
  if (sizeof(*header) != read(fd, header, sizeof(*header))) {
    int error = errno;
    JMEM_LOGV("Failed to read header for %p in %s; %s",
        addr, path, strerror(error));
    close(fd);
    unlink(path);
    return error;
  }

  // decrement and check for multiple free
  if (--header->count < 0) {

    // if we've found the multiple free, print it out
    JMEM_LOGV("multiple free %p size %d count %d",
        addr, header->size, header->count);
    char trace[2048] = {0};
    ssize_t bytes = read(fd, trace, sizeof(trace));
    if (bytes < 0) {
      int error = errno;
      JMEM_LOGV("failed to read backtrace for (%s): %s",
          path, strerror(error));
    }
   
    // print the backtraces
    JMEM_LOGV("allocation backtrace:\n"
      "%s\ndeallocation backtrace:", trace);
    uint32_t trace_size = 10;
    void *array[trace_size];
    trace_size = backtrace(array, trace_size);
    jmem_trace_record(stderr, array, trace_size);
    JMEM_LOG("end backtrace\n\n");
  }

  // we'll need to write out the header again, 
  // so seek to the beginning
  if (-1 == lseek(fd, SEEK_SET, 0)) {
    int error = errno;
    JMEM_LOGV("Failed to seek to the beginning of %s",
        path);
    close(fd);
    unlink(path);
    return error;
  }
  
  // write out the header
  if (sizeof(*header) != write(fd, header, sizeof(*header))) {
    int error = errno;
    JMEM_LOGV("Failed to update record for %p in %s: %s",
        addr, path, strerror(error));
    close(fd);
    unlink(path);
    return error;
  }

  // and all is well
  close(fd);
  return 0;
}

//*****************************************************************************
static void jmem_free(void* addr, const void* caller)
{
  // we don't care about null
  if (0 == addr) {
    return;
  }

  // restore all old hooks
  __malloc_hook = jmem_old_malloc_hook;
  __free_hook = jmem_old_free_hook;

  // do check, then free
  struct jmem_header_t header = {0};
  jmem_record_free(&header, addr);
  if (header.count >= 0) {
    free(addr);
  }

  // save hooks and restore
  jmem_old_malloc_hook = __malloc_hook;
  jmem_old_free_hook = __free_hook;
  __malloc_hook = jmem_malloc;
  __free_hook = jmem_free;
}

#endif //_jmem_h

